/**
* Copyright 2020 Huawei Technologies Co., Ltd
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at

* http://www.apache.org/licenses/LICENSE-2.0

* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.

* File sample_process.cpp
* Description: handle acl resource
*/
#include "dimp/dimp_init.hpp"
#include <iostream>
#include <string>
#include "dimp/model_process.h"
#include "acl/acl.h"
#include "utils.h"
#include <opencv2/opencv.hpp>

using namespace std;
using namespace cv;



Dimp_Init::Dimp_Init(const char* modelPath)
:deviceId_(0), imageDataBuf_(nullptr), initBoxBuf_(nullptr), transformBoxesBuf_(nullptr), isInited_(false){
    modelPath_ = modelPath;
    imageDataSize_ = RGBU8_IMAGE_SIZE(288, 288) * 13;
}

Dimp_Init::~Dimp_Init(){
    DestroyResource();
}


Result Dimp_Init::InitModel(const char* omModelPath) {

    Result ret = initmodel_.LoadModelFromFileWithMem(omModelPath);
    if (ret != SUCCESS) {
        ERROR_LOG("execute LoadModelFromFileWithMem failed");
        return FAILED;
    }

    ret = initmodel_.CreateDesc();
    if (ret != SUCCESS) {
        ERROR_LOG("execute CreateDesc failed");
        return FAILED;
    }

    ret = initmodel_.CreateOutput();
    if (ret != SUCCESS) {
        ERROR_LOG("execute CreateOutput failed");
        return FAILED;
    }


    return SUCCESS;
}

Result Dimp_Init::CreateModelInputdDataset()
{
    //Request image data memory for input model
    aclError aclRet = aclrtMalloc(&imageDataBuf_, imageDataSize_, ACL_MEM_MALLOC_HUGE_FIRST);
  //  cout << imageDataSize_<< endl;
    if (aclRet != ACL_ERROR_NONE) {
        ERROR_LOG("malloc device data buffer failed, aclRet is %d", aclRet);
        return FAILED;
    }
    //The second input to Yolov3 is the input image width and height parameter
    initBoxSize_ = sizeof(initBox_);
    transformBoxesSize_ = sizeof(transformBoxes_);

    if (runMode_ == ACL_HOST) {
        initBoxBuf_ = Utils::CopyDataHostToDevice((void *)initBox_, initBoxSize_);
        transformBoxesBuf_ = Utils::CopyDataHostToDevice((void *)transformBoxes_, transformBoxesSize_);
    }
       
    else {
        transformBoxesBuf_ = Utils::CopyDataDeviceToDevice((void *)transformBoxes_, transformBoxesSize_);
        initBoxBuf_ = Utils::CopyDataDeviceToDevice((void *)initBox_, initBoxSize_);
    }
        
    if (initBoxBuf_ == nullptr) {
        ERROR_LOG("Copy init_box info to device failed");
        return FAILED;
    }
    
    if (transformBoxesBuf_ == nullptr) {
        ERROR_LOG("Copy transform_boxes to device failed");
        return FAILED;
    }



    //Use the applied memory to create the model and input dataset. After creation, only update the memory data for each frame of inference, instead of creating the input dataset every time
    Result ret = initmodel_.Create3Inputs(imageDataBuf_, imageDataSize_,
    initBoxBuf_, initBoxSize_, transformBoxesBuf_, transformBoxesSize_);
    if (ret != SUCCESS) {
        ERROR_LOG("Create mode input dataset failed");
        return FAILED;
    }

    return SUCCESS;
}

Result Dimp_Init::Init(int32_t& deviceId, aclrtContext& context, aclrtStream& stream, aclrtRunMode& runMode) {
   if (isInited_) {
        INFO_LOG("Classify instance is initied already!");
        return SUCCESS;
    }


    deviceId_ = deviceId;
    context_ = context;
    stream_ = stream;

    runMode_ = runMode;

    aclError ret = InitModel(modelPath_);
    if (ret != SUCCESS) {
        ERROR_LOG("Init model failed");
        return FAILED;
    }

    ret = CreateModelInputdDataset();
    if (ret != SUCCESS) {
        ERROR_LOG("Create Initial Model Buffer failed!");
        return FAILED;
    }

    isInited_ = true;
    return SUCCESS;
}

Result Dimp_Init::LoadInputData(vector<cv::Mat> transform_imgs, vector<Target_Box> transform_boxes) {

    aclrtMemcpyKind policy = (runMode_ == ACL_HOST)?
                             ACL_MEMCPY_HOST_TO_DEVICE:ACL_MEMCPY_DEVICE_TO_DEVICE;

    AtlasError ret;

    uint32_t ImgSize = RGBU8_IMAGE_SIZE(288, 288);
    for (size_t i = 0; i < transform_imgs.size(); i++) {

        Mat img = transform_imgs[i].clone();
       
        ret = aclrtMemcpy(imageDataBuf_+ i * ImgSize, imageDataSize_,
                            img.data, ImgSize, policy);
        if (ret != ACL_ERROR_NONE) {
          ERROR_LOG("Copy image data to device failed.");
          return FAILED;
        }
    }

    initBox_[0] = (float) 0/1.0;
    initBox_[1] = (float)transform_boxes[0].pos.x;
    initBox_[2] = (float)transform_boxes[0].pos.y;
    initBox_[3] = (float)transform_boxes[0].pos.x + transform_boxes[0].size.x;
    initBox_[4] = (float)transform_boxes[0].pos.y + transform_boxes[0].size.y;
    // initBoxSize_ = sizeof(initBox_);
    // // cout << "initBox" <<initBox_[0]<<","<<initBox_[1]<<","<<initBox_[2]<<","<<initBox_[3]<<","<<initBox_[4]<< endl;
    // if (runMode_ == ACL_HOST)
    //     initBoxBuf_ = Utils::CopyDataHostToDevice((void *)initBox_, initBoxSize_);
    // else
    //     initBoxBuf_ = Utils::CopyDataDeviceToDevice((void *)initBox_, initBoxSize_);
    // if (initBoxBuf_ == nullptr) {
    //     ERROR_LOG("Copy init_box info to device failed");
    //     return FAILED;
    // }

    for(uint32_t i = 0; i < transform_boxes.size(); ++i)
    {
        transformBoxes_[i][0] = (float)i/1.0;
        transformBoxes_[i][1] = (float)transform_boxes[i].pos.x;
        transformBoxes_[i][2] = (float)transform_boxes[i].pos.y;
        transformBoxes_[i][3] = (float)transform_boxes[i].pos.x + transform_boxes[i].size.x;
        transformBoxes_[i][4] = (float)transform_boxes[i].pos.y + transform_boxes[i].size.y;
        // cout << "transformBoxes_" <<transformBoxes_[i][0]<<","<<transformBoxes_[i][1]<<","<<transformBoxes_[i][2]<<","<<transformBoxes_[i][3]<<","<<transformBoxes_[i][4]<< endl;
    }

    // transformBoxesSize_ = sizeof(transformBoxes_);
  
    // if (runMode_ == ACL_HOST)
    //     transformBoxesBuf_ = Utils::CopyDataHostToDevice((void *)transformBoxes_, transformBoxesSize_);
    // else
    //     transformBoxesBuf_ = Utils::CopyDataDeviceToDevice((void *)transformBoxes_, transformBoxesSize_);
    // if (transformBoxesBuf_ == nullptr) {
    //     ERROR_LOG("Copy transform_boxes to device failed");
    //     return FAILED;
    // }

    uint32_t InitBoxSize = sizeof(initBox_);
  
    ret = (Result)aclrtMemcpy(initBoxBuf_, InitBoxSize,
                            initBox_, InitBoxSize, policy);
    
    if (ret != ACL_ERROR_NONE) {
          ERROR_LOG("Copy initBox to device failed.");
          return FAILED;
        }
    
    uint32_t TransformBoxesSize = sizeof(transformBoxes_);


    ret = (Result)aclrtMemcpy(transformBoxesBuf_, TransformBoxesSize,
                            transformBoxes_, TransformBoxesSize, policy);
    
    if (ret != ACL_ERROR_NONE) {
          ERROR_LOG("Copy transformBoxes to device failed.");
          return FAILED;
        }

    return SUCCESS;
}

Result Dimp_Init::Inference() {
    aclmdlDataset* inferenceOutput = nullptr;
    //Perform reasoning
    aclError ret = initmodel_.Execute(stream_, 0);
    if (ret != SUCCESS) {
        ERROR_LOG("Execute model inference failed");
        return FAILED;
    }
    //Get inference output
    inferenceOutput = initmodel_.GetModelOutputData();

    layer2_.firstaddr = GetInferenceOutputItem(layer2_.size, inferenceOutput, 0);
    cout<<layer2_.size<<endl;

    weights_.firstaddr = GetInferenceOutputItem(weights_.size, inferenceOutput, 5);
    cout<<weights_.size<<endl;

    modulation1_.firstaddr = GetInferenceOutputItem(modulation1_.size, inferenceOutput, 3);
    cout<<modulation1_.size<<endl;

    modulation2_.firstaddr = GetInferenceOutputItem(modulation2_.size, inferenceOutput, 4);
    cout<<modulation2_.size<<endl;

    classifier_feats_.firstaddr = GetInferenceOutputItem(classifier_feats_.size, inferenceOutput, 2);
    cout<<classifier_feats_.size<<endl;

    // cout<<"******************************"<<endl;

    // float *layer2output = (float *)(layer2_.firstaddr);
    // for(int i = 0; i < 100; ++i)
    // {
    //     cout<<"layer2 = "<<layer2output[i]<<endl;
    // }
    // string modelname = "Init_Model";
    // initmodel_.OutputModelResult(modelname);

    return SUCCESS;
}


void* Dimp_Init::GetInferenceOutputItem(uint32_t& itemDataSize,
                                           aclmdlDataset* inferenceOutput,
                                           uint32_t idx) {

    aclDataBuffer* dataBuffer = aclmdlGetDatasetBuffer(inferenceOutput, idx);
    if (dataBuffer == nullptr) {
        ERROR_LOG("Get the %dth dataset buffer from model "
        "inference output failed", idx);
        return nullptr;
    }

    void* dataBufferDev = aclGetDataBufferAddr(dataBuffer);
    if (dataBufferDev == nullptr) {
        ERROR_LOG("Get the %dth dataset buffer address "
        "from model inference output failed", idx);
        return nullptr;
    }

    size_t bufferSize = aclGetDataBufferSize(dataBuffer);
    if (bufferSize == 0) {
        ERROR_LOG("The %dth dataset buffer size of "
        "model inference output is 0", idx);
        return nullptr;
    }

    void* data = nullptr;
    if (runMode_ == ACL_HOST) {
        data = Utils::CopyDataDeviceToLocal(dataBufferDev, bufferSize);
        if (data == nullptr) {
            ERROR_LOG("Copy inference output to host failed");
            return nullptr;
        }
    }
    else {
        data = dataBufferDev;
    }

    itemDataSize = bufferSize;
    return data;
}

void Dimp_Init::DestroyResource()
{

    initmodel_.Unload();
    initmodel_.DestroyDesc();
    initmodel_.DestroyInput();
    initmodel_.DestroyOutput();
}
